import torch
import torch.nn.functional as F
from tqdm import tqdm
from loss.dice_loss import dice_coeff


def eval_net(net, loader, device, batch_size):
    """Evaluation without the densecrf with the dice coefficient"""
    net.eval()
    mask_type = torch.long
    n_val = len(loader)  # the number of batch
    tot = 0

    cc = 0
    eval_corr = 0
    # with tqdm(total=n_val, desc='Validation round', unit='batch', leave=True) as pbar:
    for batch in loader:
        cc = cc + 1
        imgs, true_masks = batch['image'], batch['mask']
        imgs = imgs.to(device=device, dtype=torch.float32)
        true_masks = true_masks * 255
        true_masks = true_masks.to(device=device, dtype=mask_type)

        with torch.no_grad():
            mask_pred = net(imgs)


        true_masks = true_masks.squeeze(1)
        # tot += F.cross_entropy(mask_pred, true_masks).item()
        pred = mask_pred.max(1, keepdim=True)[1]  # get the index of the max log-probability
        eval_corr += pred.eq(true_masks.view_as(pred)).sum().item()
        # pbar.update()

    net.train()
    # print(cc)
    # return tot / n_val
    # print(eval_corr / (cc * batch_size * 512 * 512))
    return eval_corr / (cc * batch_size * 512 * 512)
